﻿CREATE or replace FUNCTION tf_nn_predict(test_table varchar,test_column varchar ,model_table varchar,result_table varchar)
  RETURNS varchar
AS $$
  import tensorflow as tf
  import numpy as np

  def predict(test_datam,model_name):
    #model_name = 'model1'
    with tf.Session() as sess:
      init = tf.global_variables_initializer()
      sess.run(init)
      #加载模型
      saver = tf.train.import_meta_graph('/plpy/model_file/%s/tf_model.meta'%model_name)
      saver.restore(sess,"/plpy/model_file/%s/tf_model"%model_name )
      graph = tf.get_default_graph()
      prediction = graph.get_tensor_by_name("czcprediction:0")
      x = graph.get_tensor_by_name("input_image:0")
      #进行预测
      pr = sess.run(prediction, feed_dict={x: test_data})
      return pr

  def insertTab(input_table,output_table,labels):
    create_sql = " create table %s as select * ,'' as label from %s where 1<>1 "%(output_table,input_table)
    plpy.execute(create_sql)
    
    insert_sql = 'insert into %s( '%output_table
    rs = plpy.execute('select * from %s order by id '%input_table)
    i = 0
    for r in rs :
      if i==0:
        for key in r :
          insert_sql = insert_sql + key +','
        insert_sql = insert_sql + 'label)values'
      insert_sql = insert_sql +"("
      for key in r :
        #如果是数组类型
        if isinstance(r[key], list):
          tstr = "{" + ','.join(str(c) for c in r[key]) + "}"
          insert_sql = insert_sql + "'" + tstr + "',"
        else :
          insert_sql = insert_sql + "'" + str(r[key]) + "',"
      insert_sql = insert_sql + "'"+ str(labels[i]) +"'),"
      i = i+1
    plpy.info(insert_sql[:-1])
    plpy.execute(insert_sql[:-1])
    
    
  # 查出待预测的数据
  q = "select %s from %s order by id " % (test_column, test_table)
  rv1 = []
  rv = plpy.execute(q)
  for i in rv :
    rv1.append(i[test_column])
  rv1 = np.array(rv1)
  test_data = np.multiply(rv1, 1.0 / 255.0)
  predict_result = predict(test_data,model_table)
  #将预测结果写入到result_table表中
  plpy.info(predict_result)
  insertTab(test_table,result_table, predict_result)
  return predict_result
$$ LANGUAGE plpythonu;


SELECT tf_nn_predict('image_tests','data','model1','result_table4');
